from tensorflow.keras.preprocessing.image import ImageDataGenerator
import tensorflow as tf
import os
from tensorflow import keras
import numpy as np

from tensorflow.compat.v1 import ConfigProto
from tensorflow.compat.v1 import InteractiveSession 

from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True

config = ConfigProto()
config.gpu_options.allow_growth = True
session = InteractiveSession(config=config)


os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

im_height = 224
im_width = 224

train_image_generator = ImageDataGenerator() 
                                             
train_data_gen = train_image_generator.flow_from_directory(directory='./brand-form_train',   # 类别文件夹 大文件夹下面是类别文件夹 可以没有图片
                                                           target_size=(im_height, im_width),
                                                           class_mode='categorical')

test_image_generator = ImageDataGenerator(rescale=1./255)

test_data_gen = test_image_generator.flow_from_directory(
	directory='./outputs',  # 分类图片文件夹 下面还要有一个子文件夹 
	target_size=(224,224),
	shuffle=False,
	batch_size=200,
	class_mode=None)

labels = (train_data_gen.class_indices)  # dict{name:index}
label = dict((v,k) for k,v in labels.items())  # dict{index:name}

model = tf.keras.models.load_model('./callbacks/mobilenetV2_2_model.h5')
pred = model.predict_generator(test_data_gen,workers=8,verbose=1)
softmax = pred.max(axis=1)  # 保存每行最大的概率值 
pred = np.argmax(pred, axis=1)  # 保存预测的是第几类

predictions = [label[index] for index in pred ]  # 转换成类名

filenames = test_data_gen.filenames
with open('real_data_predictions2.txt','w') as f:  # 分类信息保存的txt文档
	for idx in range(len(filenames)):
		str = filenames[idx] + ',' + predictions[idx] + ',' + repr(softmax[idx])
		f.write(str)
		f.write('\n')
		

